import shutil
import subprocess
import os

from typing import Optional, List, Union

from joblib import Parallel, delayed

from common.logger import logger
from common.io_utils.dirs import working_directory, in_directory


def pack_archive(
    fpath_dest: str, dir_src: str, out_format: str = "zip", remove_src: bool = False
):
    """Pack a given directory into a archival zip or tar file.

    Parameters
    ----------
    fpath_dest : str
        Destination path for the archival file.
    dir_src : str
        Source path for the directory to be packed.
    out_format : str, optional
        Output archival format, either "zip" or "tar", by default "zip".
    remove_src : bool, optional
        Whether to remove the source directory, by default False not to remove.

    Returns
    -------
    str
        Archival file path.

    Raises
    ------
    ValueError
        - Raised when the `out_format` is not one of "tar" or "zip".
        - Raised when the `fpath_dest` is under the `dir_src`.
    """

    if out_format not in ("zip", "tar"):
        raise ValueError(f"{out_format} is not supported, only support `zip` or `tar`.")

    if in_directory(fpath_dest, dir_src):
        raise ValueError(f"{fpath_dest} should not be under {dir_src}")

    with working_directory(dir_src):
        _ = shutil.make_archive(fpath_dest, out_format)

    if remove_src:
        shutil.rmtree(dir_src, ignore_errors=True)
    return fpath_dest


def zip_files(
    src_path: Union[str, List[str]],
    dest_path: str,
    zlevel: int = 7,
    verbose_level: int = 0,
    dry_run: bool = True,
) -> bool:
    """Zip a given file or list of files

    Parameters
    ----------
    src_path : Union[str, List[str]]
        A Source file path or a list of source file paths to be zipped.
    dest_path : str
        The destination zip file path.
    zlevel : int, optional
        Level of compression, in the range of 0 to 9, by default 7.
    verbose_level: int, optional
        - 0 == Do not perform compression in verbose mode (default)
        - 1+ == Perform compression in verbose mode
    dry_run : bool, optional
        Whether to run in dry-run mode, by default True.
        In dry-run mode, the operation to be run is going to printed but not actually run,
        and the function always return True.

    Returns
    -------
    bool
        Whether compression has been done successfully.
    """
    commands = ["zip", "-r", "-j", "-{}".format(zlevel)]
    if verbose_level == 0:
        commands.append("-q")
    commands.append(dest_path)

    if isinstance(src_path, list):
        for item in src_path:
            commands.append(item)
    else:
        commands.append(src_path)

    if dry_run:
        logger.info(f"dry-run {commands}")
        return True
    else:
        try:
            ret = subprocess.run(commands)
        except subprocess.CalledProcessError as err:
            logger.warning(f"fail to compress {src_path} due to {str(err)}")
            return False
        else:
            if verbose_level > 0:
                logger.info(f"finish running compression, with ret {ret}")
            if ret.returncode == 0:
                return True
            else:
                logger.warning(f"fail to compress {src_path} due to {str(ret.stderr)}")
                return False


def unzip(
    fpath_src: str,
    dir_dest: Optional[str] = None,
    exclude_files: Optional[List[str]] = None,
    overwrite: bool = False,
    verbose_level: int = 0,
):
    """Unzip a given zip file

    Parameters
    ----------
    fpath_src : str
        Path to the source zip file.
    dir_dest : str or None, optional
        Path to the destination directory. Default is None to use the current working directory.
    exclude_files : List[str] or None, optional
        List of file paths within the source zip file that are going to be excluded in unzip,
        by default None not to exclude any files.
    overwrite : bool, optional
        Whether to overwrite existing files, by default False not to overwrite.
    verbose_level: int, optional
        - 0 == do not print out process information, by default False not to print out.
        - 1+ == print out process information

    Returns
    -------
    None

    Notes
    -----
    This function directly forms the command called in a `subprocess` by using the given inputs.
    Thus please do not use it directly with external user inputs.
    """
    commands = [
        "unzip",
        "-qq" if verbose_level > 0 else "-q",
        "-o" if overwrite else "-n",
        fpath_src,
    ]
    if exclude_files is not None:
        commands.append("-x")
        commands.append(",".join(exclude_files))
    if dir_dest is not None:
        commands.append("-d")
        commands.append(dir_dest)
    subprocess.run(commands)


def unzip_batch(
    lst_fpaths: List[str],
    dir_dest: Optional[str] = None,
    nproc: Optional[int] = -1,
    overwrite: bool = False,
    verbose_level: int = 0,
):
    """Unzip a list of zip file, support parallel processing.

    Parameters
    ----------
    lst_fpaths : List[str]
        List containing paths to the source zip file.
    dir_dest : str or None, optional
        Path to the destination directory. Default is None to use the current working directory.
    nproc : int or -1, optional
        Define how many processors to use,
        by default -1 use all available processors (n_jobs=nproc=-1).
        If using a specified amount of process, nproc = a positive integer greater than 0.
        If using all but x processors, compute n = -1-x, nproc = n (e.g. all but one processor = -2).
        If bypassing joblib, nproc = 0.
    overwrite : bool, optional
        Whether to overwrite existing files, by default False not to overwrite.
    verbose_level: int, optional
        - 0 == do not print out process information, by default False not to print out.
        - 1+ == print out process information

    Returns
    -------
    None

    Notes
    -----
    This function uses the 'unzip' function,
    which directly forms the command called in a `subprocess` by using the given inputs.
    Thus please do not use it directly with external user inputs.

    > Do we need logger to record the processing?
    > we did not account for exclude file
    > consider using concurrent.futures in next revision to improve error handling

    Examples
    --------
    lst_fpaths = ['path/a.zip','path/a.zip']
    nproc = 2
    unzip_batch(lst_fpaths ,dir_dest, nproc)
    """
    if dir_dest is None:
        dir_dest = os.getcwd()
    # exclude file functionality is ignored
    exclude_files = None
    # Batch Unzipping
    if nproc == 0:
        for fpath_src in lst_fpaths:
            unzip(fpath_src, dir_dest, exclude_files, overwrite, verbose_level)
    else:
        # Define joblib process
        parallel = Parallel(n_jobs=nproc)
        parallel(
            delayed(unzip)(fpath_src, dir_dest, exclude_files, overwrite, verbose_level)
            for fpath_src in lst_fpaths
        )
